/**
 * 
 */
package mocku.transformer;

import java.lang.instrument.ClassFileTransformer;
import java.util.Arrays;
import java.util.List;
import java.util.Properties;
import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtBehavior;
import javassist.CtClass;
import javassist.CtConstructor;
import javassist.CtField;
import javassist.CtMethod;
import javassist.Modifier;
import javassist.NotFoundException;
import javassist.bytecode.AccessFlag;
import javassist.bytecode.AttributeInfo;
import javassist.bytecode.CodeAttribute;
import javassist.bytecode.LocalVariableAttribute;
import mocku.helper.StringHelper;

/**
 * Perform instrumentation on configured classes to make them act as simple
 * proxy/delegates.
 * 
 * @author aambertin
 */
public class ProxyTransformer implements ClassFileTransformer {

	Properties mockList = new Properties();

	public ProxyTransformer(Properties config) {
		mockList = config;
	}

	/**
	 * @see java.lang.instrument.ClassFileTransformer#transform(java.lang.ClassLoader,
	 *      java.lang.String, java.lang.Class, java.security.ProtectionDomain,
	 *      byte[])
	 */
	public byte[] transform(ClassLoader loader, String className, Class clazz, java.security.ProtectionDomain domain, byte[] bytes) {
		// if its in the list, then instrument!
		String cname = className.replace("/", ".");
		if (mockList.containsKey(cname)) {
			String m_conf = mockList.getProperty(cname);
			String mock = StringHelper.getMockClass(m_conf);
			String[] methods = StringHelper.getMockMethods(m_conf);
			String s_methods = "";
			for (int i = 0; i < methods.length; i++) {
				s_methods += methods[i] + " ";
			}
			System.out.println("Instrumenting " + cname + "(" + s_methods + ")");
			return doClass(cname, mock, methods, clazz, bytes);
		}

		// otherwise, just return the original bytes
		return bytes;
	}

	/*
	 * Process class for instrumentation.
	 */
	private byte[] doClass(String name, String mock, String[] methods, Class clazz, byte[] b) {
		String staticDef = "private static Class __mockClass;";
		String instanceDef = "private " + mock + " __mockInstance;";
		System.out.print("Loading pool... ");
		CtClass cl = null;
		try {
			List __methods = Arrays.asList(methods);
			ClassPool pool = ClassPool.getDefault();
			System.out.print("Loading class... ");
			cl = pool.makeClass(new java.io.ByteArrayInputStream(b));
			if (cl.isInterface() == false) {
				System.out.print("Done. Creating static field... ");
				// Give value to static field
				CtField stField = CtField.make(staticDef, cl);
				String stValue = mock + ".class;";
				cl.addField(stField, stValue);
				System.out.print("Done. Creating instance field... ");
				// Give value to mock instnace
				CtField inField = CtField.make(instanceDef, cl);
				String inValue = "new " + mock + "();";
				cl.addField(inField, inValue);

				/*
				 * Hack on constructor if methods[] contains constructor name.
				 */
				boolean interceptConstructor = __methods.isEmpty() || __methods.contains(name.substring(name.lastIndexOf(".") + 1));
					System.out.println("Done. Hacking into constructors... ");
					CtConstructor[] constructors = cl.getConstructors();
					for (int i = 0; i < constructors.length; i++) {
						if (constructors[i].isEmpty() == false) {
							doConstructor(constructors[i], mock, interceptConstructor);
						}
					}

				/*
				 * Hack on methods
				 */
				System.out.print("Done. Hacking into class methods... ");
				CtBehavior[] c_methods = cl.getDeclaredBehaviors();
				for (int i = 0; i < c_methods.length; i++) {
					if (__methods.isEmpty() || __methods.contains(c_methods[i].getName())) {
						doMethod(c_methods[i], mock);
					}
				}
				b = cl.toBytecode();
			}
		}
		catch (Exception e) {
			System.out.println("Could not instrument  " + name + ",  exception : " + e.getMessage());
			e.printStackTrace();
		}
		finally {
			if (cl != null) {
				cl.detach();
			}
		}
		// return modified bytes!
		return b;
	}

	private void doConstructor(CtConstructor constructor, String mock, boolean intercept) throws NotFoundException, CannotCompileException {
		if (!constructor.getMethodInfo().isConstructor()) return;
		
		if (intercept) {
			String body = (" { " + "this.__mockInstance = new " + mock + "(" + string_constructor_params(constructor) + ");" + " } ");
			System.out.println("New constructor body: " + body);
			constructor.setBody(body);
		}
		else {
			constructor.insertAfter("this.__mockInstance = new "+mock+"();");
		}
		constructor.insertAfter("this.__mockInstance.__reference = this; ");
	}

	public String string_constructor_params(CtConstructor constructor) throws NotFoundException {
		StringBuffer response = new StringBuffer();
		CtClass[] parTypes = constructor.getParameterTypes();
		for (int i = 0; i < parTypes.length; i++) {
			if (i > 0) response.append(",");
			response.append("$");
			response.append(i + 1);
		}
		String s_response = response.toString();
		return s_response;
	}

	/*
	 * Process methods for instrumentation.
	 */
	private void doMethod(CtBehavior method, String mock) throws NotFoundException, CannotCompileException {
		if (!method.getMethodInfo().isConstructor()) {
			System.out.print("Instrumenting method " + method.getDeclaringClass().getName() + "." + method.getName() + "... ");
			try{
				doOldWayMethod(method, mock);
			}catch (Throwable e){
				System.out.print("First try to instrument method fail (" + e.getMessage() + "), trying other way");
				doNewWayMathod(method, mock);
			}
			System.out.print("DONE.");
		}
	}
	
	private void doOldWayMethod(CtBehavior method, String mock) throws CannotCompileException, NotFoundException{
		if(AccessFlag.isPrivate(method.getModifiers())){
			System.out.print("Method is private, setting public... ");
			method.setModifiers(AccessFlag.setPublic(method.getModifiers()));
		}
		
		if(Modifier.isStatic(method.getModifiers())){
			System.out.print("Method is static... ");
		}
		
		method.setBody(" return __mockInstance." + method.getName() + "(" + string_params(method) + "); ");
	}

	private void doNewWayMathod(CtBehavior method, String mock) throws CannotCompileException, NotFoundException{
		if (method instanceof CtMethod) {
			CtMethod met = ((CtMethod)method);
			CtClass[] methodParams = met.getParameterTypes();

			ClassPool pool = ClassPool.getDefault();
			CtClass mockClass = pool.get(mock);

			CtMethod newMethod = mockClass.getDeclaredMethod(method.getName(), methodParams);

			if(newMethod != null){
				((CtMethod)method).setBody(newMethod, null);

			} else {
				System.out.print("Could not find method with same parameters in mock... ");
			}
		} else {
			System.out.print("Error instrumenting method... ");
		}
	}
	
	/**
	 * return String for paramTypes definition.
	 * 
	 * @param method
	 * @return
	 * @throws NotFoundException
	 */
	private String string_param_types(CtBehavior method) throws NotFoundException {
		StringBuffer params = new StringBuffer(" new Class[] { ");
		boolean first = true;
		for (CtClass type : method.getParameterTypes()) {
			if (first) {
				first = false;
			}
			else {
				params.append(",");
			}
			params.append(type.getName());
		}
		params.append("}");
		return params.toString();
	}

	/**
	 * Return string with parameters used to make proxy calls.
	 * 
	 * @param method
	 * @return
	 * @throws NotFoundException
	 */
	private String string_params(CtBehavior method) throws NotFoundException {
		CtClass[] types = method.getParameterTypes();
		CodeAttribute codeAttribute = method.getMethodInfo().getCodeAttribute();

		LocalVariableAttribute locals = null;
		StringBuffer paramsList = new StringBuffer(" ");

		if (codeAttribute != null) {
			AttributeInfo attribute;
			attribute = codeAttribute.getAttribute("LocalVariableTable");
			locals = (LocalVariableAttribute) attribute;
		}
		for (int i = 0; i < types.length; i++) {
			if (i > 0) {
				paramsList.append(",");
			}
			//paramsList.append(parameterNameFor(method, locals, i));
			paramsList.append("$"+(i+1));
		}

		paramsList.append(" ");

		return paramsList.toString();
	}

	/**
	 * Return parameter names for method, used to makes proxy calls.
	 * 
	 * @param method
	 * @param locals
	 * @param i
	 * @return
	 */
	private String parameterNameFor(CtBehavior method, LocalVariableAttribute locals, int i) {

		if (locals == null) {
			return Integer.toString(i + 1);
		}
		int modifiers = method.getModifiers();
		int j = i;
		if (Modifier.isSynchronized(modifiers)) {
			j++;
		}
		if (Modifier.isStatic(modifiers) == false) {
			j++;
		}
		String variableName = locals.variableName(j);
		return variableName;
	}

}
